import os
import ast

def get_class_name_from_file(file_path):
    """从 .py 文件中提取第一个类名"""
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            file_content = f.read()

        tree = ast.parse(file_content)
        for node in ast.walk(tree):
            if isinstance(node, ast.ClassDef):
                return node.name  # 返回第一个类名
    except Exception as e:
        print(f"⚠️ 无法解析 {file_path}：{e}")
    return None


def rename_files_in_folder(folder_path):
    """递归遍历文件夹并重命名 .py 文件（自动处理重名）"""
    used_names = {}  # 记录每个类名出现次数：{类名: 次数}

    for root, _, files in os.walk(folder_path):
        for file_name in files:
            if file_name.endswith(".py") and file_name != "__init__.py":
                file_path = os.path.join(root, file_name)
                class_name = get_class_name_from_file(file_path)

                if not class_name:
                    print(f"⚠️ 未找到类定义，跳过：{file_name}")
                    continue

                # 检查重名次数
                used_names[class_name] = used_names.get(class_name, 0) + 1
                if used_names[class_name] > 1:
                    new_file_name = f"{class_name}_{used_names[class_name]}.py"
                else:
                    new_file_name = f"{class_name}.py"

                new_file_path = os.path.join(root, new_file_name)

                # 如果文件名已相同则跳过
                if file_name == new_file_name:
                    continue

                # 若已存在同名文件，继续增加序号避免覆盖
                count = used_names[class_name]
                while os.path.exists(new_file_path):
                    count += 1
                    new_file_name = f"{class_name}_{count}.py"
                    new_file_path = os.path.join(root, new_file_name)
                    used_names[class_name] = count

                os.rename(file_path, new_file_path)
                print(f"✅ 重命名：{file_name} → {new_file_name}")

    print("\n🎯 所有文件处理完成！")


if __name__ == "__main__":
    folder = input("请输入要遍历的文件夹路径：").strip()
    if not os.path.isdir(folder):
        print("❌ 输入的路径不是有效文件夹。")
    else:
        rename_files_in_folder(folder)
